package com.encode.huffman;

import java.util.*;

/**
 * @Author lichunlong lcl253137@alibaba-inc.com
 * @Desc 霍夫曼编码
 * @Date 2025-03-27
 */
public class HuffmanCode {

    /**
     * 赫夫曼编码表{32=01, 97=100, 100=11000, 117=11001, 101=1110, 118=11011, 105=101, 121=11010, 106=0010, 107=1111, 108=000, 111=0011}
     */
    public static final Map<Byte, String> HUFFMAN_CODE_TABLE_MAP = new HashMap<>();

    /**
     * 赫夫曼编码路径: 左子树
     */
    public static final String HUFFMAN_CODE_LEFT_PATH = "LEFT_PATH";

    /**
     * 赫夫曼编码路径: 右子树
     */
    public static final String HUFFMAN_CODE_RIGHT_PATH = "RIGHT_PATH";

    /**
     * 赫夫曼编码路径MAP
     * 左子树路径为0，右子树路径为1
     */
    public static final Map<String, String> HUFFMAN_CODE_PATH_MAP = new HashMap<>() {{
        put(HUFFMAN_CODE_LEFT_PATH, "0");
        put(HUFFMAN_CODE_RIGHT_PATH, "1");
    }};

    /**
     * 将Byte字节数组转换成霍夫曼节点集合
     *
     * @param byteArray 字节数组
     * @return 霍夫曼节点集合[Node[date= 97, weight= 5], Node[date= 32,weight= 9]], ...]
     */
    public static List<HuffmanNode> initNodeList(byte[] byteArray) {
        List<HuffmanNode> nodeList = new ArrayList<>();

        // 遍历byte数组, 统计每一个byte出现的次数 -> map[byteDataChar, occurTime]
        Map<Byte, Integer> occurTimeMap = new HashMap<>();
        for (byte byteData : byteArray) {
            occurTimeMap.compute(byteData, (k, occurTime) -> Objects.isNull(occurTime) ? 1 : occurTime + 1);
        }

        for (Map.Entry<Byte, Integer> entry : occurTimeMap.entrySet()) {
            nodeList.add(new HuffmanNode(entry.getKey(), entry.getValue()));
        }

        return nodeList;
    }

    /**
     * 创建赫夫曼树
     *
     * @param nodeList 霍夫曼节点集合
     * @return 根节点
     * @see com.neusoft.tree.huffman.HuffmanTree#createHuffmanTree(int[])
     */
    public static HuffmanNode createHuffmanTree(List<HuffmanNode> nodeList) {
        while (nodeList.size() > 1) {
            // 1、处理前先排序
            Collections.sort(nodeList);

            // 2、取出根节点权值最小的二叉树的2个Node节点
            HuffmanNode subLeftNode = nodeList.get(0);
            HuffmanNode subRightNode = nodeList.get(1);

            // 创建一颗新的二叉树,它的根节点没有data, 只有权值
            HuffmanNode parentNode = new HuffmanNode(null, subLeftNode.getWeight() + subRightNode.getWeight(), subLeftNode, subRightNode);

            nodeList.remove(subLeftNode);
            nodeList.remove(subRightNode);

            nodeList.add(parentNode);
        }

        return nodeList.get(0);
    }

    /**
     * 前序遍历霍夫曼树
     *
     * @param HuffmanRootNode 霍夫曼树的根节点
     */
    public static void preOrder(HuffmanNode HuffmanRootNode) {
        if (Objects.isNull(HuffmanRootNode)) {
            System.out.println("赫夫曼树为空");
            return;
        }

        HuffmanRootNode.preOrder();
    }

    /**
     * 初始化霍夫曼编码表
     *
     * @param HuffmanRootNode 霍夫曼树的根节点
     */
    public static void initHuffmanCodeTable(HuffmanNode HuffmanRootNode) {
        if (Objects.isNull(HuffmanRootNode)) {
            return;
        }

        StringBuilder tempStringBuilder = new StringBuilder();
        // 处理根结点的左子树
        initHuffmanCodeTable(HuffmanRootNode.getLeftSubNode(), HUFFMAN_CODE_PATH_MAP.get(HUFFMAN_CODE_LEFT_PATH), tempStringBuilder);

        //处理根结点的右子树
        initHuffmanCodeTable(HuffmanRootNode.getRightSubNode(), HUFFMAN_CODE_PATH_MAP.get(HUFFMAN_CODE_RIGHT_PATH), tempStringBuilder);
    }

    /**
     * 初始化霍夫曼编码表
     *
     * @param node                霍夫曼树的节点
     * @param codePath            当前节点的路径
     * @param parentStringBuilder 当前节点的父节点路径
     */
    public static void initHuffmanCodeTable(HuffmanNode node, String codePath, StringBuilder parentStringBuilder) {
        StringBuilder currentStringBuilder = new StringBuilder(parentStringBuilder);
        currentStringBuilder.append(codePath);

        if (Objects.isNull(node)) {
            return;
        }
        if (Objects.nonNull(node.getData())) {
            // 找到了叶子结点: 叶子结点肯定是原始节点
            HUFFMAN_CODE_TABLE_MAP.put(node.getData(), currentStringBuilder.toString());
            return;
        }

        // 向左递归
        initHuffmanCodeTable(node.getLeftSubNode(), HUFFMAN_CODE_PATH_MAP.get(HUFFMAN_CODE_LEFT_PATH), currentStringBuilder);
        // 向右递归
        initHuffmanCodeTable(node.getRightSubNode(), HUFFMAN_CODE_PATH_MAP.get(HUFFMAN_CODE_RIGHT_PATH), currentStringBuilder);
    }

    /**
     * 霍夫曼编码
     *
     * @param byteArray 字符串对应的字节数组
     * @return 霍夫曼编码后的字节数组
     */
    public static byte[] encode(byte[] byteArray) {
        /*
         * 1、利用霍夫曼编码表将字节数组转成赫夫曼编码对应的字符串
         * "i like like like java do you like a java" => "1010100010111111110010001011111111001000101111111100100101001101110001110000011011101000111100101000101111111100110001001010011011100"
         */
        StringBuilder stringBuilder = new StringBuilder();
        for (byte byteData : byteArray) {
            stringBuilder.append(HUFFMAN_CODE_TABLE_MAP.get(byteData));
        }

        // 2、将"1010100010111111110..."转成byte[]字节数组: 8位对应一个byte
        int length = stringBuilder.length();
        // 统计返回的byte数组长度
        int byteArrayLength = length % 8 == 0 ? length / 8 : length / 8 + 1;
        byte[] huffmanCodeByteArray = new byte[byteArrayLength];
        int index = 0;

        for (int i = 0; i < stringBuilder.length(); i += 8) {
            String byteString = (i + 8 > stringBuilder.length() ? stringBuilder.substring(i) : stringBuilder.substring(i, i + 8));
            huffmanCodeByteArray[index] = (byte) Integer.parseInt(byteString, 2);
            index++;
        }

        return huffmanCodeByteArray;
    }

    /**
     * 霍夫曼解码为二进制的字符串
     *
     * @param byteData 霍夫曼编码后的字节
     * @param flag     是否需要补高位标志; true表示需要补高位; false表示不需要补, 最后一个字节无需补高位
     * @return 二进制的字符串
     */
    public static String decode2BinaryString(byte byteData, Boolean flag) {
        int tempIntData = byteData;
        if (flag) {
            // 按位与256: 1 0000 0000 | 0000 0001 => 1 0000 0001
            tempIntData |= 256;
        }
        // 返回tempIntData对应二进制的补码
        String binaryString = Integer.toBinaryString(tempIntData);

        return flag ? binaryString.substring(binaryString.length() - 8) : binaryString;
    }

    /**
     * 霍夫曼解码
     *
     * @param huffmanCodeByteArray 霍夫曼编码后的字节数组
     * @return 解码后的字节数组
     */
    public static byte[] decode(byte[] huffmanCodeByteArray) {
        // 1、先得到huffmanCodeByteArray对应的二进制的字符串
        StringBuilder stringBuilder = new StringBuilder();
        for (int index = 0; index < huffmanCodeByteArray.length; index++) {
            byte huffmanCodeByteData = huffmanCodeByteArray[index];
            // 判断是不是最后一个字节
            boolean flag = (index == huffmanCodeByteArray.length - 1);
            stringBuilder.append(decode2BinaryString(huffmanCodeByteData, !flag));
        }

        // 2、反转赫夫曼编码表: 因为反向查询 a-> 100; 100-> a
        Map<String, Byte> huffmanCodeTableReverseMap = new HashMap<String, Byte>();
        for (Map.Entry<Byte, String> entry : HUFFMAN_CODE_TABLE_MAP.entrySet()) {
            huffmanCodeTableReverseMap.put(entry.getValue(), entry.getKey());
        }

        // 3、把二进制的字符串按照翻转的赫夫曼编码表进行解码
        List<Byte> byteList = new ArrayList<>();
        for (int sbIndex = 0; sbIndex < stringBuilder.length(); ) {
            int counter = 1;
            boolean flag = true;
            Byte byteData = null;

            while (flag) {
                // sbIndex不动, 让counter移动, 直到匹配到一个字符
                String reverseMapKey = stringBuilder.substring(sbIndex, sbIndex + counter);
                byteData = huffmanCodeTableReverseMap.get(reverseMapKey);
                if (Objects.isNull(byteData)) {
                    // 没有匹配到
                    counter++;
                } else {
                    // 匹配到, 退出循环
                    flag = false;
                }
            }
            byteList.add(byteData);
            // sbIndex直接移动到counter
            sbIndex += counter;
        }

        // 4、将byteList转换成byte[]并返回
        byte[] byteArray = new byte[byteList.size()];
        for (int i = 0; i < byteArray.length; i++) {
            byteArray[i] = byteList.get(i);
        }
        return byteArray;
    }
}
